# -*- coding=UTF-8 -*-
import sys
import os
import random
import cv2
import math
import numpy as np
import tensorflow as tf
import linecache
import string

TRAIN_DATA_DIR = '../train_data/'
TEST_DATA_DIR = '../test_data/'

class batchData:
	trainImage = []
	trainLabel = []
	
	testImage = []
	testLabel = []
	
	def __init__(self):
		# train data ----------------------- #
		if not os.path.exists(TRAIN_DATA_DIR):
			print 'error: ', TRAIN_DATA_DIR, 'not exists'
			return
		fileList = os.listdir(TRAIN_DATA_DIR)
		print '>> infor: ', TRAIN_DATA_DIR, 'exists, files number', len(fileList)
		for curFile in fileList:
			fg = curFile.find('.jpg')
			if fg < 0:
				continue
			fg = curFile.rfind('_')
			if fg < 0:
				continue
			subStr0 = curFile[fg + 1 : len(curFile)]
			fg = subStr0.find('.')
			if fg < 0:
				continue
			subStr1 = subStr0[0 : fg]
			pathImg = TRAIN_DATA_DIR + curFile
			num = string.atoi(subStr1, base = 10)
			
			curImage = cv2.imread(pathImg, flags = cv2.IMREAD_GRAYSCALE)
			curImage.astype('float32') / 255.0
			curLabel = [0.,0.,0.,0.,0.,0.,0.,0.,0.,0.]
			curLabel[num] = 1
			self.trainImage.append(curImage)
			self.trainLabel.append(curLabel)
			#invert image
			ivtImage = 255 - curImage
			ivtImage.astype('float32') / 255.0
			self.trainImage.append(ivtImage)
			self.trainLabel.append(curLabel)
		print '>> infor: trainImage number', len(self.trainImage)
		# test data ----------------------- #
		if not os.path.exists(TEST_DATA_DIR):
			print '>> infor: ', TEST_DATA_DIR, 'not exists'
			return
		fileList = os.listdir(TEST_DATA_DIR)
		print '>> infor: ', TEST_DATA_DIR, 'exists, files number', len(fileList)
		for curFile in fileList:
			fg = curFile.find('.jpg')
			if fg < 0:
				continue
			fg = curFile.rfind('_')
			if fg < 0:
				continue
			subStr0 = curFile[fg + 1 : len(curFile)]
			fg = subStr0.find('.')
			if fg < 0:
				continue
			subStr1 = subStr0[0 : fg]
			pathImg = TEST_DATA_DIR + curFile
			num = string.atoi(subStr1, base = 10)
			
			curImage = cv2.imread(pathImg, flags = cv2.IMREAD_GRAYSCALE)
			curImage.astype('float32') / 255.0
			curLabel = [0.,0.,0.,0.,0.,0.,0.,0.,0.,0.]
			curLabel[num] = 1
			self.testImage.append(curImage)
			self.testLabel.append(curLabel)
		print '>> infor: testImage number', len(self.testImage)
		print '>> infor: init OK'
	
	def next_batch_train(self, batchSize = 200):
		# 打乱顺序，每次取打乱顺序后的前batchSize个
		order = range(len(self.trainImage))
		random.shuffle(order)
		curBatchImage = []
		curBatchLabel = []
		for i in range(0, batchSize):
			curBatchImage.append(self.trainImage[order[i]])
			curBatchLabel.append(self.trainLabel[order[i]])
		return np.expand_dims(np.array(curBatchImage), -1), np.array(curBatchLabel)

	def all_test(self):
		return np.expand_dims(np.array(self.testImage), -1), np.array(self.testLabel)

if __name__ == "__main__":
	print '--STA--'
	
	dataset = batchData()
	curBatchImage, curBatchLabel = dataset.next_batch()
	print curBatchImage.shape, curBatchLabel.shape
	print '--END--'
